from A import *
import os
import numpy as np

# find temperature drift data and imu_tk data
# input param
temp_data_dirs = "../test_imu_error_data/temp_drift_data"  # 温漂测试文件路径
imu_tk_data_dirs = "../test_imu_error_data/imu_tk_data"    # imu_tk测试文件路径
temp_drift_dirs = "../test_imu_error_data/drift_cof"       # 温漂保存数据路径
data_hz = 500                                              # 数据采样频率
data_long = 3000                                           # 温漂数据最大拟合长度
run_name = "/home/zz/imu_tk/bin/test_imu_calib"            # imu_tk 可执行文件路径
save_dir = "../test_imu_error_data/imu_tk_cof"             # imu_tk 标定文件存储路径

# imu_tk 临时文件存储路径
dir_ab_comp = os.path.abspath("./")
dir_data_comp = "temp_data"
data_save_name = os.path.join(dir_ab_comp, dir_data_comp)

if not os.path.exists(data_save_name):
    os.system("mkdir" + " " + data_save_name)

# process temp_drift_data
dir_ab, lists = find_file(temp_data_dirs)
for i in range(len(lists)):
    acc_data, gyro_data = pre_process(os.path.join(dir_ab,lists[i]), data_hz, data_long)
    acc_cof, gyro_cof = temperature_fitting(acc_data, gyro_data, data_hz)

    acc_cof_name = os.path.join(temp_drift_dirs,"acc_cof_"+lists[i][-7:-4]+".txt")
    gyro_cof_name = os.path.join(temp_drift_dirs, "gyro_cof_" + lists[i][-7:-4] + ".txt")

    np.savetxt(acc_cof_name, acc_cof)
    np.savetxt(gyro_cof_name, gyro_cof)

# pre process imu_tk_data
dir_ab, lists = find_file(imu_tk_data_dirs)
print(dir_ab)
print(lists)
for i in range(len(lists)):
    acc_data, gyro_data = pre_process(os.path.join(dir_ab, lists[i]), data_hz, 3000)

    acc_cof_name = os.path.join(temp_drift_dirs, "acc_cof_" + lists[i][-7:-4] + ".txt")
    acc_cof_data = np.loadtxt(acc_cof_name)
    print("acc_cof_data")
    print(acc_cof_data)
    print(acc_cof_name)
    acc_data[:, 1] = acc_data[:, 1] - (acc_cof_data[0, 0] * acc_data[:, 4] + acc_cof_data[0, 1])
    acc_data[:, 2] = acc_data[:, 2] - (acc_cof_data[1, 0] * acc_data[:, 4] + acc_cof_data[1, 1])
    acc_data[:, 3] = acc_data[:, 3] - (acc_cof_data[2, 0] * acc_data[:, 4] + acc_cof_data[2, 1])

    gyro_cof_name = os.path.join(temp_drift_dirs, "gyro_cof_" + lists[i][-7:-4] + ".txt")
    gyro_cof_data = np.loadtxt(gyro_cof_name)
    print(gyro_cof_name)
    print(gyro_cof_data)
    gyro_data[:, 1] = gyro_data[:, 1] - (gyro_cof_data[0, 0] * gyro_data[:, 4] + gyro_cof_data[0, 1])
    gyro_data[:, 2] = gyro_data[:, 2] - (gyro_cof_data[1, 0] * gyro_data[:, 4] + gyro_cof_data[1, 1])
    gyro_data[:, 3] = gyro_data[:, 3] - (gyro_cof_data[2, 0] * gyro_data[:, 4] + gyro_cof_data[2, 1])

    np.savetxt(data_save_name + "/" + "acc_data_" + lists[i][-7:-4] + ".txt", acc_data)
    np.savetxt(data_save_name + "/" + "gyro_data_" + lists[i][-7:-4] + ".txt", gyro_data)

    param1 = data_save_name + "/" + "acc_data_" + lists[i][-7:-4] + ".txt"
    param2 = data_save_name + "/" + "gyro_data_" + lists[i][-7:-4] + ".txt"
    param3 = os.path.abspath(save_dir) + "/" + "acc_cof_" + lists[i][-7:-4] + ".txt"
    param4 = os.path.abspath(save_dir) + "/" + "gyro_cof_" + lists[i][-7:-4] + ".txt"

    rp = run_name + " " + param1 + " " + param2 + " " + param3 + " " + param4
    os.system(rp)

os.system("rm" + " " + data_save_name + "/" + " " + "-r")
